"""
来自论文[1] 许利恒.基于椭圆特征的陨石坑检测与识别方法研究[D].北京:北京航空航天大学,2022

按照论文的方法复现，关于从四点金字塔拓展至整幅陨石坑的算法，这里使用了姿态解算的结果作为重投影平差的依据，不再使用导航坑对的不变量计算，增加了对陨石坑的筛选算法
"""

import numpy as np
from ...base import Matching
from utils.ellipse import center_ellipse, radius_ellipse
from itertools import combinations
from scipy.optimize import linear_sum_assignment
from utils.pose import pose_calculate
import cv2
import pandas as pd


class TriadHash(Matching):
    def __init__(self, catalog_dir, **kwargs):
        super().__init__(catalog_dir, **kwargs)
        self.descriptor = pd.Series(self.descriptor.item())

    def triad_descriptor(
        self, i, j, k, C1, C2, C3, U1, U2, U3, th, *args, **kwargs
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray[bool]]:
        """
        从陨石坑的参数中计算描述子，C1和C2可以是含有不确定度的椭圆参数
        TODO 必须支持向量化操作
        Output:
          这里的输出一定是编码后的Hash值！
        """
        raise NotImplementedError

    def identify(self, params, uncertainty, factor, *args, max_try=5, **kwargs):
        """
        输入应该越接近于真实的检测输出越好。因此输入应当是一堆椭圆的拟合参数，输出应当是匹配确定的结果，即椭圆ID的列表。
        Arguments:
            params (np.ndarray) : ellipse parameters for (x^2, xy, y^2, x, y, 1)
        Returns:
            list : matched ellipse ID
        """
        # 将陨石坑按直径排序
        params = np.array(params)
        diameters = np.mean(radius_ellipse(params), axis=0)
        idx = np.argsort(-diameters)
        if len(idx) < 4:
            return None
        fail_num = 0
        for i in range(idx.shape[0]):
            j, k, l = np.array(list(combinations(np.delete(idx, i), 3))).T
            i = idx[i].repeat(k.shape[0])
            # 批量算出金字塔的三元组
            qs, ijk, index, valid = self.triad_descriptor(
                np.concatenate((i, i, j, i)),
                np.concatenate((j, j, k, k)),
                np.concatenate((k, l, l, l)),
                params[np.concatenate((i, i, j, i))],
                params[np.concatenate((j, j, k, k))],
                params[np.concatenate((k, l, l, l))],
                uncertainty[np.concatenate((i, i, j, i))],
                uncertainty[np.concatenate((j, j, k, k))],
                uncertainty[np.concatenate((k, l, l, l))],
                th=2,
            )
            IJK, ind = self.hash_search(qs)
            # 变成四个三元组，即金字塔的结构
            valid = (valid & ind).reshape(4, -1).all(axis=0)
            if not valid.any():
                fail_num += 1
                continue
            IJK = IJK.reshape(4, -1)[..., valid]
            ijk = ijk.reshape(3, 4, -1)[..., valid]
            index = index.reshape(3, 4, -1)[..., valid]
            for it in range(ijk.shape[2]):
                pyra = ijk[..., it]
                assert np.unique(pyra[2, 1:]).shape[0] == 1
                ijkl = np.array((*pyra[:, 0], pyra[2, 1]))
                # 对冗余三元组还原为原始顺序，与ijk对应
                IJK_pyra = [
                    np.take_along_axis(IJK[i, it].T, index[:, i, None, it], axis=0)
                    for i in range(4)
                ]
                for IJKL in self.build_pyramids(*IJK_pyra):
                    success, result = self.extend_identify(
                        params, uncertainty, ijkl, IJKL, *args, factor=factor, **kwargs
                    )
                    if success:
                        return result
            fail_num += 1
            if fail_num > max_try:
                break
        return None

    def hash_search(self, qs):
        """
        输入是金字塔的哈希值，输出也应当是金字塔的哈希值吗？
        """
        # 取消投票法，改为使用金字塔法
        v = self.descriptor.reindex(qs.flatten())
        v = pd.DataFrame(v.values.reshape(qs.shape[0], -1)).apply(
            lambda row: row.dropna().tolist(), axis=0
        )
        # 将按行有元素的合并为同一个array
        ind = v.apply(len) > 0
        v[ind] = v[ind].apply(lambda row: np.concatenate(row, axis=0))
        # 按行收在一起，以取得至少有一个冗余三元组的列
        return v.to_numpy(), ind.to_numpy()

    def build_pyramids(self, IJK, IJL, JKL, IKL):
        """以IJK为底座"""
        L_IJ_ind = (IJK[0, :, None] == IJL[0]) & (IJK[1, :, None] == IJL[1])
        L_JK_ind = (IJK[1, :, None] == JKL[0]) & (IJK[2, :, None] == JKL[1])
        L_IK_ind = (IJK[0, :, None] == IKL[0]) & (IJK[2, :, None] == IKL[1])
        num_ind = np.argwhere(
            L_IJ_ind.any(axis=1) & L_JK_ind.any(axis=1) & L_IK_ind.any(axis=1)
        ).squeeze(1)
        if num_ind.shape[0] == 0:
            return []
        # 加强条件，只取各行均为一个元素的行
        ind = (
            (L_IJ_ind[num_ind].sum(axis=1) == 1)
            & (L_JK_ind[num_ind].sum(axis=1) == 1)
            & (L_IK_ind[num_ind].sum(axis=1) == 1)
        )
        L_IJ = IJL[2, None].repeat(num_ind.shape[0], axis=0)[
            L_IJ_ind[num_ind] & ind[:, None]
        ]
        L_JK = JKL[2, None].repeat(num_ind.shape[0], axis=0)[
            L_JK_ind[num_ind] & ind[:, None]
        ]
        L_IK = IKL[2, None].repeat(num_ind.shape[0], axis=0)[
            L_IK_ind[num_ind] & ind[:, None]
        ]
        Ls_ind = (L_IJ == L_JK) & (L_IJ == L_IK)
        # 恢复原始顺序
        return np.stack((*IJK[:, num_ind[ind][Ls_ind]], L_IJ[Ls_ind]), axis=1)

    def extend_identify(
        self,
        params,
        uncertainty,
        ijkl,
        IJKL,
        K,
        th: float,
        confidence: float,
        dist_th=5,
        **kwargs,
    ):
        """
        TODO 用增量式重投影方法完成对全幅陨石坑的匹配，具体的细节如下：
        1. 利用现有四点ijkl估计出一个摄像机矩阵P（TODO 有可能估计失败）
        2. 根据当前的四点ijkl，设置某个超参数r，搜索四点周围r个点构成候选重投影点，利用P重投影至图像平面，记为Qr
        3. 利用匈牙利算法，对当前的Qr+4个重投影点与当前图像上的全部陨石坑中心点（可能多于Qr+4）作匹配（TODO 算法输出差过大的点应当舍去）。
        4. 对于每个匹配上的陨石坑中心点，计算其不变量以排除潜在的噪声点。
        4. 匹配的结果作为新的ijkl，重新估计P，重复步骤2-3，直到全部匹配上或者达到最大迭代次数
        """
        # 返回第一符合要求的值
        # 将金字塔识别结果拓展至整个图像
        if len(ijkl) < 4:
            return False, None
        cnts = np.array(center_ellipse(params)).T
        ijkl, IJKL = self.increment_hungary(params, ijkl, IJKL, K, dist_th)
        if len(ijkl) < 4:
            return False, None
            # 重投影验证
        points_3d = self.catalog[IJKL, :]
        obse_2d = cnts[ijkl, :]
        R, T = pose_calculate(K, points_3d, obse_2d)
        if R is None:
            return False, None
        uv = cv2.projectPoints(points_3d, R, T, K, np.zeros(4))[0].squeeze()
        ind = np.linalg.norm(uv - obse_2d, axis=1) < th
        if ind.mean() >= confidence:
            return True, (ijkl[ind], IJKL[ind])
        else:
            return False, None

    def increment_hungary(
        self,
        params,
        uncertainty,
        ijkl: tuple,
        IJKL: tuple,
        K,
        dist_th=2.0,
    ):
        """
        用增量式重投影方法完成对全幅陨石坑的匹配，具体的细节如下：
        1. 利用现有四点ijkl估计出一个摄像机矩阵P（TODO 有可能估计失败）
        2. 根据当前的四点ijkl，设置某个超参数r，搜索四点周围r个点构成候选重投影点，利用P重投影至图像平面，记为Qr
        3. 利用匈牙利算法，对当前的Qr+4个重投影点与当前图像上的全部陨石坑中心点（可能多于Qr+4）作匹配；
        4. 对于每个匹配上的陨石坑中心点，计算其不变量以排除潜在的噪声点。
        4. 匹配的结果作为新的ijkl，重新估计P，重复步骤2-3，直到全部匹配上或者达到最大迭代次数
        Arguments:
            ijkl (tuple) : 当前图像每个被匹配的陨石坑序号
            IJKL (tuple) : 当前图像每个被匹配的陨石坑在目录中的序号
            K (np.ndarray) : 相机内参矩阵
            dist_th (float) : 匹配距离阈值
        """
        # 估计P
        points_3d = self.catalog[IJKL, :]
        cnts = np.array(center_ellipse(params)).T
        obse_2d = cnts[ijkl, :]
        R, T = pose_calculate(K, points_3d, obse_2d)
        if R is None:
            return [], []
        # 重投影
        uv = cv2.projectPoints(self.catalog, R, T, K, np.zeros(4))[0].squeeze()
        # 把当前可能位于图像平面内的全部点都加入至候选列表中
        ind = (
            (uv[:, 0] < K[0, 2] * 2)
            & (uv[:, 0] >= 0)
            & (uv[:, 1] < K[1, 2] * 2)
            & (uv[:, 1] >= 0)
        )
        # 得到全部候选中心点列表
        cata_3d = uv[ind]
        cata_id = np.arange(self.catalog.shape[0])[ind]
        # 获取当前全部图像中心点
        obse_id = np.arange(len(cnts))
        # 匈牙利算法
        dist_matrix = np.linalg.norm(cata_3d[:, np.newaxis] - cnts, axis=2)
        # 使用匈牙利算法计算最优分配
        row_ind, col_ind = linear_sum_assignment(dist_matrix)
        # 最优分配
        cata_id = cata_id[row_ind]
        obse_id = obse_id[col_ind]
        # 计算匹配误差，本误差也将作为筛选真实陨石坑和虚假陨石坑的依据
        error = dist_matrix[row_ind, col_ind]
        # 计算不变量验证
        ## 具体的步骤应当包括：
        ## 1. 根据匈牙利算法匹配出的距离远近和设计的超参数dist_th，选出符合匹配距离范围内的候选陨石坑
        ## 2. 全部候选陨石坑两两作对，算出不变量以及配对
        ## 3. 根据配对关系，给出导航坑表内的不变量对
        # 挑选匹配距离范围内的陨石坑
        ind = error < dist_th
        cata_id = cata_id[ind]
        obse_id = obse_id[ind]
        return obse_id, cata_id
